/*
 * Wazuh Vulnerability scanner - Scan Orchestrator
 * Copyright (C) 2015, Wazuh Inc.
 * March 13, 2024.
 *
 * This program is free software; you can redistribute it
 * and/or modify it under the terms of the GNU General Public
 * License (version 2) as published by the FSF - Free Software
 * Foundation.
 */

#ifndef _SCAN_AGENT_LIST_HPP
#define _SCAN_AGENT_LIST_HPP

#include "agentReScanListException.hpp"
#include "chainOfResponsability.hpp"
#include "loggerHelper.h"
#include "scanContext.hpp"
#include "socketDBWrapper.hpp"
#include "wazuhDBQueryBuilder.hpp"
#include "wdbDataException.hpp"

/**
 * @brief Orchestrates queries to the global Wazuh-DB system and initiates package scanning.
 *
 * The `TScanAgentList` class is responsible for coordinating the execution of queries within the
 * broader Wazuh ecosystem.
 *
 * @tparam TScanContext scan context type.
 */
template<typename TScanContext = ScanContext,
         typename TAbstractHandler = AbstractHandler<std::shared_ptr<TScanContext>>,
         typename TSocketDBWrapper = SocketDBWrapper>
class TScanAgentList final : public AbstractHandler<std::shared_ptr<TScanContext>>
{
private:
    std::shared_ptr<TAbstractHandler> m_packageScanSuborchestration;
    std::shared_ptr<TAbstractHandler> m_osScanSuborchestration;

    /**
     * @brief The `scanAgentOs` method fetches the OS information of an agent from the Wazuh-DB and sends the data to
     * the `osScanSuborchestration` component for further processing.
     *
     * @param agent The agent data.
     */
    void scanAgentOs(const AgentData& agent, const bool noIndex)
    {
        nlohmann::json response;

        try
        {
            TSocketDBWrapper::instance().query(WazuhDBQueryBuilder::builder().agentGetOsInfoCommand(agent.id).build(),
                                               response);
        }
        catch (const SocketDbWrapperException& e)
        {
            throw WdbDataException(e.what(), agent.id);
        }
        catch (std::exception& e)
        {
            logError(WM_VULNSCAN_LOGTAG,
                     "Unable to retrieve OS agent data (agent %s). Reason: %s.",
                     agent.id.c_str(),
                     e.what());
            return;
        }

        // Validate the response
        if (response.empty())
        {
            logDebug2(WM_VULNSCAN_LOGTAG, "Empty response for agent '%s' in Wazuh-DB 'sys_' query", agent.id.c_str());
            return;
        }

        for (const auto& osInfo : response)
        {
            flatbuffers::FlatBufferBuilder fbBuilder;
            auto agentInfo = Synchronization::CreateAgentInfoDirect(
                fbBuilder, agent.id.c_str(), agent.ip.c_str(), agent.name.c_str(), agent.version.c_str());

            auto syscollectorOs =
                Synchronization::Createsyscollector_osinfoDirect(fbBuilder,
                                                                 osInfo.value("architecture", "").c_str(),
                                                                 osInfo.value("checksum", "").c_str(),
                                                                 osInfo.value("hostname", "").c_str(),
                                                                 osInfo.value("os_build", "").c_str(),
                                                                 osInfo.value("os_codename", "").c_str(),
                                                                 osInfo.value("os_display_version", "").c_str(),
                                                                 osInfo.value("os_major", "").c_str(),
                                                                 osInfo.value("os_minor", "").c_str(),
                                                                 osInfo.value("os_name", "").c_str(),
                                                                 osInfo.value("os_patch", "").c_str(),
                                                                 osInfo.value("os_platform", "").c_str(),
                                                                 osInfo.value("os_release", "").c_str(),
                                                                 osInfo.value("os_version", "").c_str(),
                                                                 osInfo.value("release", "").c_str(),
                                                                 osInfo.value("scan_time", "").c_str(),
                                                                 osInfo.value("sysname", "").c_str(),
                                                                 osInfo.value("version", "").c_str());

            auto stateMsg =
                Synchronization::Createstate(fbBuilder,
                                             Synchronization::AttributesUnion::AttributesUnion_syscollector_osinfo,
                                             syscollectorOs.Union());

            auto syncMsg = Synchronization::CreateSyncMsg(
                fbBuilder, agentInfo, Synchronization::DataUnion::DataUnion_state, stateMsg.Union());

            fbBuilder.Finish(syncMsg);

            std::variant<const SyscollectorDeltas::Delta*, const Synchronization::SyncMsg*, const nlohmann::json*>
                variantData = Synchronization::GetSyncMsg(fbBuilder.GetBufferPointer());

            auto context = std::make_shared<TScanContext>(variantData);
            context->m_noIndex = noIndex;

            m_osScanSuborchestration->handleRequest(std::move(context));
        }
    }

    /**
     * @brief The `scanAgentPackages` method fetches the package information of an agent from the Wazuh-DB and sends the
     * data to the `packageScanSuborchestration` component for further processing.
     *
     * @param agent The agent data.
     */
    void scanAgentPackages(const AgentData& agent, const bool noIndex)
    {
        nlohmann::json response;

        try
        {
            TSocketDBWrapper::instance().query(WazuhDBQueryBuilder::builder().agentGetPackagesCommand(agent.id).build(),
                                               response);
        }
        catch (const SocketDbWrapperException& e)
        {
            throw WdbDataException(e.what(), agent.id);
        }
        catch (std::exception& e)
        {
            logError(WM_VULNSCAN_LOGTAG,
                     "Unable to retrieve packages agent data (agent %s). Reason: %s.",
                     agent.id.c_str(),
                     e.what());
            return;
        }

        // Validate the response
        if (response.empty())
        {
            logDebug2(
                WM_VULNSCAN_LOGTAG, "Empty response for agent '%s' in Wazuh-DB 'sys_programs' query", agent.id.c_str());
            return;
        }

        for (const auto& package : response)
        {
            flatbuffers::FlatBufferBuilder fbBuilder;
            auto agentInfo = Synchronization::CreateAgentInfoDirect(
                fbBuilder, agent.id.c_str(), agent.ip.c_str(), agent.name.c_str(), agent.version.c_str());
            auto syscollectorPackage =
                Synchronization::Createsyscollector_packagesDirect(fbBuilder,
                                                                   package.value("architecture", "").c_str(),
                                                                   package.value("checksum", "").c_str(),
                                                                   package.value("description", "").c_str(),
                                                                   package.value("format", "").c_str(),
                                                                   package.value("groups", "").c_str(),
                                                                   package.value("install_time", "").c_str(),
                                                                   package.value("item_id", "").c_str(),
                                                                   package.value("location", "").c_str(),
                                                                   package.value("multiarch", "").c_str(),
                                                                   package.value("name", "").c_str(),
                                                                   package.value("priority", "").c_str(),
                                                                   package.value("scan_time", "").c_str(),
                                                                   package.value("size", 0),
                                                                   package.value("source", "").c_str(),
                                                                   package.value("vendor", "").c_str(),
                                                                   package.value("version", "").c_str());

            auto stateMsg =
                Synchronization::Createstate(fbBuilder,
                                             Synchronization::AttributesUnion::AttributesUnion_syscollector_packages,
                                             syscollectorPackage.Union());

            auto syncMsg = Synchronization::CreateSyncMsg(
                fbBuilder, agentInfo, Synchronization::DataUnion::DataUnion_state, stateMsg.Union());

            fbBuilder.Finish(syncMsg);

            std::variant<const SyscollectorDeltas::Delta*, const Synchronization::SyncMsg*, const nlohmann::json*>
                variantData = Synchronization::GetSyncMsg(fbBuilder.GetBufferPointer());

            auto context = std::make_shared<TScanContext>(variantData);
            context->m_noIndex = noIndex;

            m_packageScanSuborchestration->handleRequest(std::move(context));
        }
    }

public:
    /**
     * @brief Construct a new `TScanAgentList` object.
     *
     * The constructor instantiates a new `TScanAgentList` object with the provided package and OS scan orchestration
     * instances.
     *
     * @param packageScanOrchestration package orchestration instance.
     * @param osScanOrchestration os orchestration instance.
     */
    explicit TScanAgentList(std::shared_ptr<TAbstractHandler> packageScanOrchestration,
                            std::shared_ptr<TAbstractHandler> osScanOrchestration)
        : m_packageScanSuborchestration(std::move(packageScanOrchestration))
        , m_osScanSuborchestration(std::move(osScanOrchestration))
    {
    }

    /**
     * @brief Handles request and passes control to the next step of the chain.

     * The `handleRequest` method processes incoming requests, executes queries in the Wazuh-DB, and forwards the
     * query response to a sub-orchestration component. The handling involves several steps:
     *
     * 1. It iterates over the list of agents and executes queries to fetch data from the Wazuh-DB.
     * 2. For each agent, it scans the OS and packages.
     * 3. If a query fails, it logs a warning and adds the agent to a list of agents with incomplete scans.
     * 4. If the list is not empty, it throws an exception to initiate a rescan for the agents in the list.
     * @param data A shared pointer to the input data.
     * @return A shared pointer to the result of the request processing.
     *
     */
    std::shared_ptr<TScanContext> handleRequest(std::shared_ptr<TScanContext> data) override
    {
        for (const auto& agent : data->m_agents)
        {
            try
            {
                scanAgentOs(agent, data->m_noIndex);
                scanAgentPackages(agent, data->m_noIndex);
            }
            catch (const WdbDataException& e)
            {
                logDebug2(
                    WM_VULNSCAN_LOGTAG, "Error executing query to fetch agent data for agents. Reason: %s.", e.what());
                data->m_agentsWithIncompletedScan.push_back(agent);
            }
            catch (const std::exception& e)
            {
                logError(WM_VULNSCAN_LOGTAG, "Error handling request: %s.", e.what());
            }
        }

        if (!data->m_agentsWithIncompletedScan.empty())
        {
            throw AgentReScanListException(
                "Error executing rescan for multiple agents.", data->m_agentsWithIncompletedScan, data->m_noIndex);
        }
        return AbstractHandler<std::shared_ptr<TScanContext>>::handleRequest(std::move(data));
    }
};

using ScanAgentList = TScanAgentList<>;

#endif // _SCAN_AGENT_LIST_HPP
